package onem.soria.utils;

import lombok.SneakyThrows;
import onem.soria.entity.constants.Constant;

import javax.crypto.Cipher;
import javax.crypto.spec.SecretKeySpec;
import java.util.Base64;

public class DecryptUtil {
    /**
     * RC4-KSA algorithm to generate the key box.
     * @param key the key to use for generating the key box
     * @return the generated key box
     */
    public static byte[] generateKeyBox(byte[] key) {
        byte[] keyBox = new byte[Constant.RC4_LENGTH];
        for (int i = 0; i < Constant.RC4_LENGTH; i++) {
            keyBox[i] = (byte) i;
        }
        for (int i = 0, j = 0; i < Constant.RC4_LENGTH; i++) {
            j = (j + keyBox[i] + key[i % key.length]) & Constant.RC4_CONSTANT;
            byte swap = keyBox[i];
            keyBox[i] = keyBox[j];
            keyBox[j] = swap;
        }
        return keyBox;
    }

    /**
     * RC4-PRGA algorithm to generate the key stream.
     * @param text the text to encrypt
     * @param keyBox the key box to use for generating the key stream
     * @return the generated key stream
     */
    public static byte[] generateKeyStream(byte[] text, byte[] keyBox) {
        byte[] keyStream = new byte[Constant.RC4_LENGTH];
        for (int i = 0; i < Constant.RC4_LENGTH; i++) {
            int index = (i + keyBox[i]) & 255;
            int index2 = keyBox[i] + keyBox[index];
            keyStream[i] = keyBox[index2 & Constant.RC4_CONSTANT];
        }
        for (int i = 0; i < text.length; i++) {
            text[i] ^= keyStream[(i + 1) % Constant.RC4_LENGTH];
        }
        return text;
    }

    /**
     * Decrypts a base64 encoded byte array using the provided key.
     * @param text the base64 encoded byte array to decrypt
     * @return the decrypted byte array
     */
    public static byte[] decryptBase64(byte[] text) {
        Base64.Decoder decoder = Base64.getDecoder();
        return decoder.decode(text);
    }

    /**
     * Decrypts a byte array using the provided key with AES algorithm.
     * @param text the byte array to decrypt
     * @param key the key to use for decryption
     * @return the decrypted byte array
     */
    @SneakyThrows
    public static byte[] decryptCore(byte[] text, byte[] key) {
        SecretKeySpec specification = new SecretKeySpec(key, Constant.DECRYPT_CORE_ALGORITHM);
        Cipher cipher = Cipher.getInstance(Constant.DECRYPT_CORE_MODE);
        cipher.init(Cipher.DECRYPT_MODE, specification);
        return cipher.doFinal(text);
    }
}
